#include <iostream>
#include "util/distance.hpp"
#include "util/time_memory.hpp"
#include <vector>
#include <cstring>
#include <cmath>
#include <unordered_set>
#include <queue>
#include <random>

#undef NDEBUG

#include <cassert>
#include <algorithm>

/*
 * 假设当前迭代点为cand, 邻居为n, 查询为q
 * 终止条件:
 *      向量cand q在cand n上的投影与向量cand n的模进行比较, 如果小就全部搜索, 否则只搜索梯度边
 */

namespace GraphGradient {

    class GraphGradient {
    public:
        GraphGradient(int vec_dim, int max_element, int random_seed = 100) {
            max_element_ = max_element;
            vec_dim_ = vec_dim;
            srand(random_seed);
            report_every_ = 100000;
        }

        GraphGradient() {

        }

        void add(float *data, std::vector<std::vector<int>> &graph, std::vector<std::vector<double>> &graph_distance,
                 int enter_point);

        std::vector<std::unordered_set<int>> search_knn(float *query, int k, int n_query, int ef_search);

        std::vector<double> get_gradient_proportion_l(double &avg_n_gradient, double &avg_n_neighbors);

        ~GraphGradient() {
        }

        std::vector<std::vector<int>> graph_;
        std::vector<int> gradient_size_l_;
        std::vector<double> dist_to_ep_l_;
        std::vector<std::vector<double>> dist_graph_;

#ifdef TEST
        bool is_warm_up_ = true;
        unsigned int test_avg_n_hops_ = 0;
        unsigned int test_avg_n_visit_edges_ = 0;
        float test_avg_n_gradient_hops_ = 0;
        unsigned int test_avg_n_visited_edges_ = 0;

        std::vector<unsigned int> test_avg_n_hops_l;
        std::vector<unsigned int> test_avg_n_visit_edges_l;
        std::vector<float> test_avg_n_gradient_hops_l;
        std::vector<unsigned int> test_avg_n_visited_edges_l;

        int gradient_size0_ = 0;
        int neighbor_size0_ = 0;
#endif

#ifdef PRINT_ANGLE
        unsigned int test_avg_ep_ep_neighbor_obtuse_edges_ = 0;
        unsigned int test_avg_ep_ep_query_obtuse_edges_ = 0;
        unsigned int test_avg_ep_neighbor_query_obtuse_edges_ = 0;
        unsigned int test_avg_ep_acute_edges_ = 0;
        unsigned int test_avg_cand_cand_neighbor_obtuse_edges_ = 0;
        unsigned int test_avg_cand_cand_query_obtuse_edges_ = 0;
        unsigned int test_avg_cand_neighbor_query_obtuse_edges_ = 0;
        unsigned int test_avg_cand_acute_edges_ = 0;

        std::vector<unsigned int> test_avg_ep_ep_neighbor_obtuse_edges_l;
        std::vector<unsigned int> test_avg_ep_ep_query_obtuse_edges_l;
        std::vector<unsigned int> test_avg_ep_neighbor_query_obtuse_edges_l;
        std::vector<unsigned int> test_avg_ep_acute_edges_l;
        std::vector<unsigned int> test_avg_cand_cand_neighbor_obtuse_edges_l;
        std::vector<unsigned int> test_avg_cand_cand_query_obtuse_edges_l;
        std::vector<unsigned int> test_avg_cand_neighbor_query_obtuse_edges_l;
        std::vector<unsigned int> test_avg_cand_acute_edges_l;
#endif

#ifdef TEST_PRINT
        std::vector<std::vector<std::vector<int>>> test_search_answer_path_l_l_l;
        std::vector<std::vector<int>> test_search_answer_path_l_l;
        std::vector<std::vector<std::vector<double>>> test_path_distance_l_l_l;
        std::vector<std::vector<double>> test_path_distance_l_l;
#endif

    private:
        float *data_;
        int max_element_;
        int vec_dim_;
        int enterpoint_;
        size_t report_every_;

        std::unordered_set<int> search_labeled(float *item, int k, int ef_search);

    };

    /*
     * use the distance to enter point as the judgement of gradient and neighbor
     */
    void GraphGradient::add(float *data, std::vector<std::vector<int>> &graph,
                            std::vector<std::vector<double>> &graph_distance, int enter_point) {
        data_ = data;

        enterpoint_ = enter_point;

        float *ep_data = data_ + enterpoint_ * vec_dim_;
        dist_to_ep_l_.resize(max_element_);
        for (int i = 0; i < max_element_; i++) {
            dist_to_ep_l_[i] = l2distance(data + i * vec_dim_, ep_data, vec_dim_);
        }

        gradient_size_l_.resize(max_element_);
        graph_.resize(max_element_);
        dist_graph_.resize(max_element_);
        TimeRecord single_insert_record;
        for (int i = 0; i < max_element_; i++) {
            std::vector<int> graph_i = graph[i];
            int graph_i_size = graph_i.size();
            graph_[i].resize(graph_i_size);
            dist_graph_[i].resize(graph_i_size);
            int gradient_size = 0;
            int start_idx = 0;
            int end_idx = graph_i_size - 1;
            for (std::vector<int>::iterator it = graph_i.begin(); it != graph_i.end(); ++it) {
                if (dist_to_ep_l_[i] < dist_to_ep_l_[*it] || i == enterpoint_) {
                    dist_graph_[i][start_idx] = graph_distance[i][*it];
                    graph_[i][start_idx] = *it;
                    start_idx++;
                    gradient_size++;
                } else {
                    dist_graph_[i][end_idx] = graph_distance[i][*it];
                    graph_[i][end_idx] = *it;
                    end_idx--;
                }
            }
            gradient_size_l_[i] = gradient_size;

            if (i % report_every_ == 0) {
                std::cout << i / (0.01 * max_element_) << " %, "
                          << 1e-6 * single_insert_record.get_elapsed_time_micro() << " s/iter" << " Mem: "
                          << get_current_RSS() / 1000000 << " Mb \n";
                single_insert_record.reset();
            }
        }

#ifdef TEST
        gradient_size0_ = 0;
        neighbor_size0_ = 0;
        for (int i = 0; i < max_element_; i++) {
            int gradient_size = gradient_size_l_[i];
            int neighbor_size = graph_[i].size() - gradient_size;
            gradient_size0_ = gradient_size == 0 ? gradient_size0_ + 1 : gradient_size0_;
            neighbor_size0_ = neighbor_size == 0 ? neighbor_size0_ + 1 : neighbor_size0_;
        }
        std::printf("gradient_size0 %d, neighbor_size0 %d\n", gradient_size0_, neighbor_size0_);
#endif
    }


    /*
     * TODO 尝试着将std::vector<std::unordered_set<int>>提到外面去
     */
    std::vector<std::unordered_set<int>> GraphGradient::search_knn(float *query, int k, int n_query, int ef_search) {
        std::vector<std::unordered_set<int>> res(n_query);
        float *tmp_query = query;
        for (int i = 0; i < n_query; i++, tmp_query += vec_dim_) {
            res[i] = search_labeled(tmp_query, k, ef_search);
        }
#ifdef TEST
        test_avg_n_hops_ = test_avg_n_hops_ / n_query;
        test_avg_n_visit_edges_ = test_avg_n_visit_edges_ / n_query;
        test_avg_n_gradient_hops_ = test_avg_n_gradient_hops_ / n_query;
        test_avg_n_visited_edges_ = test_avg_n_visited_edges_ / n_query;

        if(!is_warm_up_){
            test_avg_n_hops_l.push_back(test_avg_n_hops_);
            test_avg_n_visit_edges_l.push_back(test_avg_n_visit_edges_);
            test_avg_n_gradient_hops_l.push_back(test_avg_n_gradient_hops_);
            test_avg_n_visited_edges_l.push_back(test_avg_n_visited_edges_);
        }

        test_avg_n_hops_ = 0;
        test_avg_n_visit_edges_ = 0;
        test_avg_n_gradient_hops_ = 0;
        test_avg_n_visited_edges_ = 0;



#endif
#ifdef PRINT_ANGLE
        test_avg_ep_ep_neighbor_obtuse_edges_ = test_avg_ep_ep_neighbor_obtuse_edges_ / n_query;
        test_avg_ep_ep_query_obtuse_edges_ = test_avg_ep_ep_query_obtuse_edges_ / n_query;
        test_avg_ep_neighbor_query_obtuse_edges_ = test_avg_ep_neighbor_query_obtuse_edges_ / n_query;
        test_avg_ep_acute_edges_ = test_avg_ep_acute_edges_ / n_query;
        test_avg_cand_cand_neighbor_obtuse_edges_ = test_avg_cand_cand_neighbor_obtuse_edges_ / n_query;
        test_avg_cand_cand_query_obtuse_edges_ = test_avg_cand_cand_query_obtuse_edges_ / n_query;
        test_avg_cand_neighbor_query_obtuse_edges_ = test_avg_cand_neighbor_query_obtuse_edges_ / n_query;
        test_avg_cand_acute_edges_ = test_avg_cand_acute_edges_ / n_query;

        if(!is_warm_up_){
            test_avg_ep_ep_neighbor_obtuse_edges_l.push_back(test_avg_ep_ep_neighbor_obtuse_edges_);
            test_avg_ep_ep_query_obtuse_edges_l.push_back(test_avg_ep_ep_query_obtuse_edges_);
            test_avg_ep_neighbor_query_obtuse_edges_l.push_back(test_avg_ep_neighbor_query_obtuse_edges_);
            test_avg_ep_acute_edges_l.push_back(test_avg_ep_acute_edges_);
            test_avg_cand_cand_neighbor_obtuse_edges_l.push_back(test_avg_cand_cand_neighbor_obtuse_edges_);
            test_avg_cand_cand_query_obtuse_edges_l.push_back(test_avg_cand_cand_query_obtuse_edges_);
            test_avg_cand_neighbor_query_obtuse_edges_l.push_back(test_avg_cand_neighbor_query_obtuse_edges_);
            test_avg_cand_acute_edges_l.push_back(test_avg_cand_acute_edges_);
        }

        test_avg_ep_ep_neighbor_obtuse_edges_ = 0;
        test_avg_ep_ep_query_obtuse_edges_ = 0;
        test_avg_ep_neighbor_query_obtuse_edges_ = 0;
        test_avg_ep_acute_edges_ = 0;
        test_avg_cand_cand_neighbor_obtuse_edges_ = 0;
        test_avg_cand_cand_query_obtuse_edges_ = 0;
        test_avg_cand_neighbor_query_obtuse_edges_ = 0;
        test_avg_cand_acute_edges_ = 0;
#endif
#ifdef TEST_PRINT
        test_search_answer_path_l_l_l.push_back(test_search_answer_path_l_l);
        test_search_answer_path_l_l = std::vector<std::vector<int>>();
        test_path_distance_l_l_l.push_back(test_path_distance_l_l);
        test_path_distance_l_l = std::vector<std::vector<double>>();
#endif

        return res;
    }

    //返回的数组是按照相对query距离排序的, 降序, 就是第一个元素是数组中离query最远的元素
    std::unordered_set<int> GraphGradient::search_labeled(float *item, int k, int ef_search) {

#ifdef TEST_PRINT
        std::vector<int> test_search_answer_path_l;
        std::vector<double> test_path_distance_l;
#endif

        std::unordered_set<int> visited_element;
        std::priority_queue<edge_stop_condition, std::vector<edge_stop_condition>, edge_stop_condition_min_heap> candidates;
        std::priority_queue<distance_cmp, std::vector<distance_cmp>, distance_max_heap> nearest_neighbors;
        double dist_ep_query = l2distance(data_ + enterpoint_ * vec_dim_, item, vec_dim_);
        /*
         * sift configuration
         * double dist_ep_q_1 = dist_ep_query * 0.8;
        double dist_ep_q_2 = dist_ep_query * 1.2;
         */
        double dist_ep_q_1 = dist_ep_query * 0.4;
        double dist_ep_q_2 = dist_ep_query * 1.4;
        {
            visited_element.insert(enterpoint_);
            candidates.emplace(dist_ep_query, enterpoint_, 0, dist_ep_query);
            nearest_neighbors.emplace(dist_ep_query, enterpoint_);
        }
        while (candidates.size() > 0) {
            edge_stop_condition cand = candidates.top();
#ifdef TEST
            test_avg_n_hops_++;
#endif
#ifdef TEST_PRINT
            test_search_answer_path_l.push_back(cand.idx_);
            test_path_distance_l.push_back(cand.query_dist_);
#endif
            candidates.pop();
            distance_cmp far_ele = nearest_neighbors.top();
            if (cand.query_dist_ > far_ele.distance_) {
                break;
            }

            double ep_neighbor_dist = dist_to_ep_l_[cand.idx_];
#ifdef PRINT_ANGLE

            if(dist_ep_query + cand.query_dist_ <= ep_neighbor_dist){
                test_avg_ep_ep_neighbor_obtuse_edges_++;
            }else if(ep_neighbor_dist + cand.query_dist_ <= dist_ep_query){
                test_avg_ep_ep_query_obtuse_edges_++;
            }else if(ep_neighbor_dist + dist_ep_query <= cand.query_dist_){
                test_avg_ep_neighbor_query_obtuse_edges_++;
            }
            if(dist_ep_query + cand.query_dist_ > ep_neighbor_dist
                        && ep_neighbor_dist + cand.query_dist_ > dist_ep_query
                        && ep_neighbor_dist + dist_ep_query > cand.query_dist_){
                test_avg_ep_acute_edges_++;
            }

            if(cand.query_dist_ + cand.last_query_dist_ <= cand.last_this_dist_){
                test_avg_cand_cand_neighbor_obtuse_edges_++;
            }else if(cand.query_dist_ + cand.last_this_dist_ <= cand.last_query_dist_){
                test_avg_cand_cand_query_obtuse_edges_++;
            }else if(cand.last_query_dist_ + cand.last_this_dist_ <= cand.query_dist_){
                test_avg_cand_neighbor_query_obtuse_edges_++;
            }

            if(cand.query_dist_ + cand.last_query_dist_ > cand.last_this_dist_
                    && cand.query_dist_ + cand.last_this_dist_ > cand.last_query_dist_
                    && cand.last_query_dist_ + cand.last_this_dist_ > cand.query_dist_){
                test_avg_cand_acute_edges_++;
            }
#endif
            bool is_gradient = false;

            std::vector<int> neighbor_l = graph_[cand.idx_];
            int end_idx = neighbor_l.size();
            if (is_gradient) {
#ifdef TEST
                test_avg_n_gradient_hops_++;
#endif
                end_idx = gradient_size_l_[cand.idx_];
            }
            std::vector<double> cand_neighbor_dist_l = dist_graph_[cand.idx_];
            for (int i = 0; i < end_idx; i++) {
                int neighbor = neighbor_l[i];
                if (visited_element.find(neighbor) == visited_element.end()) { //can not find
#ifdef TEST
                    test_avg_n_visit_edges_++;
#endif
                    visited_element.insert(neighbor);
                    far_ele = nearest_neighbors.top();
                    double neighbor_query_dist = l2distance(item, data_ + neighbor * vec_dim_, vec_dim_);

                    if ((int) nearest_neighbors.size() < ef_search || neighbor_query_dist < far_ele.distance_) {

                        candidates.emplace(neighbor_query_dist, neighbor,
                                           cand_neighbor_dist_l[i], cand.query_dist_);
                        nearest_neighbors.emplace(neighbor_query_dist, neighbor);
                        while ((int) nearest_neighbors.size() > ef_search) {
                            nearest_neighbors.pop();
                        }
                    }


                }
#ifdef TEST
                else{
                    test_avg_n_visited_edges_++;
                }
#endif
            }

        }
#ifdef TEST_PRINT
        test_search_answer_path_l_l.push_back(test_search_answer_path_l);
        test_path_distance_l_l.push_back(test_path_distance_l);
#endif
        while ((int) nearest_neighbors.size() > k) {
            nearest_neighbors.pop();
        }

        std::unordered_set<int> res;
        while (nearest_neighbors.size() > 0) {
            distance_cmp tmp = nearest_neighbors.top();
            nearest_neighbors.pop();
            res.insert(tmp.idx_);
        }

        return res;
    }

    std::vector<double> GraphGradient::get_gradient_proportion_l(double &avg_n_gradient, double &avg_n_neighbors) {
        avg_n_gradient = 0;
        avg_n_neighbors = 0;
        std::vector<double> graph_proportion(max_element_);
        for (int i = 0; i < max_element_; i++) {
            int neighbors_size = graph_[i].size();
            int gradient_size = gradient_size_l_[i];
            avg_n_gradient += gradient_size;
            avg_n_neighbors += neighbors_size;
            double proportion = (double) gradient_size / (double) neighbors_size;
            graph_proportion[i] = proportion;
        }
        avg_n_gradient = avg_n_gradient / max_element_;
        avg_n_neighbors = avg_n_neighbors / max_element_;
        return graph_proportion;
    }

}